import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

from ultralytics import YOLO

def main():
    # 获取当前目录的绝对路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 加载预训练模型
    model = YOLO('yolov8n-pose.pt')
    
    # 开始训练
    results = model.train(
        data=os.path.join(current_dir, 'dataset.yaml'),  # 数据集配置文件
        epochs=150,                   # 训练轮数
        imgsz=640,                   # 图像大小
        batch=4,                     # 批次大小
        device=None,                 # 自动选择设备
        patience=30,                 # 早停耐心值
        project='pickleball_pose',   # 项目名称
        name='train1',              # 运行名称
        exist_ok=True,              # 允许覆盖已存在的实验目录
        pretrained=True,            # 使用预训练权重
        optimizer='Adam',           # 优化器
        verbose=True,              # 显示详细信息
        plots=True,                # 绘制训练图表
        # 基本数据增强
        hsv_h=0.015,              # 图像HSV-Hue增强
        hsv_s=0.7,                # 图像HSV-Saturation增强
        hsv_v=0.4,                # 图像HSV-Value增强
        degrees=5.0,              # 图像旋转角度
        translate=0.1,            # 图像平移
        scale=0.5,                # 图像缩放
        shear=2.0                 # 图像剪切
    )

if __name__ == "__main__":
    main() 